import time
import os
import pickle
import numpy as np
import faiss
from typing import List, Dict
from embed import get_embedding
from config import FAISS_INDEX_PATH, VECTOR_DIM


class EntertainmentRAGStore:
    def __init__(self):
        self.index_path = FAISS_INDEX_PATH
        self.index_file = os.path.join(self.index_path, "faiss_rag.index")
        self.meta_file = os.path.join(self.index_path, "metadata.pkl")

        # IVF算法参数（根据数据量调整）
        self.nlist = 100  # 聚类数量，数据量越大需要越大（如100-10000）

        # 创建目录
        os.makedirs(self.index_path, exist_ok=True)

        if os.path.exists(self.index_file):
            # 加载已存在的索引
            self.index = faiss.read_index(self.index_file)
        else:
            # 初始化IVF索引（替代原来的FlatL2）
            # 量化器：使用FlatL2作为聚类中心的计算方式
            quantizer = faiss.IndexFlatL2(VECTOR_DIM)
            # 创建IVF索引：量化器 + 向量维度 + 聚类数量
            self.index = faiss.IndexIVFFlat(quantizer, VECTOR_DIM, self.nlist)
            # 新索引初始为未训练状态

        if os.path.exists(self.meta_file):
            with open(self.meta_file, 'rb') as f:
                self.metadata = pickle.load(f)
        else:
            self.metadata = []

    def add_texts(self, texts: List[str]):
        """批量添加文本到IVF索引（包含训练步骤）"""
        vectors = []
        new_meta = []

        for text in texts:
            vec = get_embedding(text)
            vectors.append(vec)
            new_meta.append({"text": text, "added_at": time.time()})

        # 转为numpy数组（FAISS要求的格式）
        vectors_np = np.array(vectors, dtype=np.float32)

        # IVF索引必须先训练才能添加数据（首次添加时训练）
        if not self.index.is_trained:
            # 用当前向量作为训练数据（需保证数量 > nlist，否则会警告）
            self.index.train(vectors_np)
            print(f"✅ IVF索引训练完成（聚类数：{self.nlist}）")

        # 添加向量到索引
        self.index.add(vectors_np)
        self.metadata.extend(new_meta)

        # 持久化索引和元数据
        faiss.write_index(self.index, self.index_file)
        with open(self.meta_file, 'wb') as f:
            pickle.dump(self.metadata, f)

        print(f"✅ 已添加 {len(texts)} 条文本到IVF索引（总条数：{len(self)}）")

    def __len__(self):
        return self.index.ntotal


if __name__ == '__main__':
    store = EntertainmentRAGStore()
    # 测试添加数据（首次运行会自动训练索引）
    store.add_texts(['aaaaabbbbb', '测试IVF索引', '这是一条娱乐新闻'])